# -*- coding: utf-8 -*-

import torch
from torch.utils.data import Dataset
from torchvision.datasets.folder import default_loader
from pathlib import Path

from feature_extraction import resnet_transform
import h5py
import numpy as np


def load_h5(data_dir, data_name, label_type):
    """
    label_type = 1: use gt_1(score) as labels;
    label_type = 2: use gt_2(bin) as labels.
    """
    feature = []
    label = []
    weight = []  # TODO how to use weight
    filename = data_dir + 'Data_' + data_name + '_google_p5.h5'
    f = h5py.File(filename)
    # flatten video index into a 1D array, and sort it in ascending order
    # NOTE: 特例
    # np.sort(np.array(dataset['youtube']['ord']).astype('int32').flatten())
    #   ->  array([11, 12, ..., 21, 23, ..., 50], dtype=int32)
    video_ord = np.sort(np.array(f['ord']).astype('int32').flatten())
    for i in video_ord:  # for each video
        feature.append(np.matrix(f['fea_' + i.__str__()]).astype('float32'))
        label.append(np.array(f[f'gt_{label_type}_{i}']).astype('float32'))
        weight.append(np.array(label_type - 1.0).astype('float32'))
    f.close()

    return feature, label, weight


class VideoData(Dataset):
    def __init__(self, preprocessed=True,
                 transform=resnet_transform, with_name=False,
                 data=None, label=None, vdo_id=None
                 ):
        # TODO root, is_train等并未使用, 历史原因
        # self.root = root
        self.preprocessed = preprocessed
        self.transform = transform
        self.with_name = with_name
        self.data = data
        self.label = label
        self.vdo_id = vdo_id
        # self.is_train = is_train

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        if self.preprocessed:
            # image_path = self.video_list[index]
            # with h5py.File(image_path, 'r') as f:
            # NOTE: (Pdb) print(list(self.dataset['youtube'].keys()))

            if self.with_name:    # TODO 假设with_name时为测试
                # 1024, 707
                # return torch.Tensor(np.array(f['pool5'])), image_path.name[:-5]
                return torch.Tensor(np.array(self.data[index])), f'fea_{self.vdo_id[index]+1}', self.label[index]
            else:
                # 1024, 707
                return torch.Tensor(np.array(self.data[index])), self.label[index]
                # return torch.Tensor(np.array(f['pool5']))

        else:
            pass
            # images = []
            # for img_path in Path(self.video_list[index]).glob('*.jpg'):
            #     img = default_loader(img_path)
            #     img_tensor = self.transform(img)
            #     images.append(img_tensor)
            #
            # return torch.stack(images), img_path.parent.name[4:]


if __name__ == '__main__':
    pass
